#include <fstream>
#include <iostream>
#include <sstream>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>
}

void saveFrame(AVFrame *frame, int width, int height, int frameIndex) {
  std::stringstream filename;
  filename << "frame" << frameIndex << ".ppm";
  std::ofstream file(filename.str(), std::ios_base::binary);

  // header
  file << "P6\n" << width << " " << height << "\n255\n";

  // pixel data
  for (int i = 0; i < height; ++i) {
    file.write((const char *)(frame->data[0] + i * frame->linesize[0]),
               width * 3);
  }
}

// Making Screencaps
int main(int argc, char *argv[]) {
  int ret(0);
  AVFormatContext *formatCtx(nullptr);
  AVCodecContext *codecCtx(nullptr);
  AVFrame *frame(nullptr);
  AVFrame *frameRGB(nullptr);
  uint8_t *buffer(nullptr);

  // 取命令行参数获取文件名
  std::string filename;
  if (argc > 1) {
    filename = argv[1];
  } else {
    filename = "C:/Users/zhou/Videos/bigbuckbunny.mp4";
  }

  // 打开文件获取上下文
  ret = avformat_open_input(&formatCtx, filename.c_str(), nullptr, nullptr);
  if (ret < 0) {
    std::cout << "error: avformat_open_input() failed" << std::endl;
    goto fail;
  }

  // 初始化流信息
  ret = avformat_find_stream_info(formatCtx, nullptr);
  if (ret < 0) {
    std::cout << "error: avformat_find_stream_info() failed" << std::endl;
    goto fail;
  }

  // 打印媒体文件信息
  av_dump_format(formatCtx, 0, filename.c_str(), 0);

  // 查找视频流
  int videoStream(-1);
  for (uint32_t i = 0; i < formatCtx->nb_streams; ++i) {
    if (formatCtx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
      videoStream = i;
      break;
    }
  }

  if (videoStream == -1) {
    std::cout << "error: Didn't find a video stream";
    goto fail;
  }

  // 获取解码器
  const AVCodec *codec =
      avcodec_find_decoder(formatCtx->streams[videoStream]->codecpar->codec_id);
  if (!codec) {
    std::cout << "error: Unsupported codec" << std::endl;
    goto fail;
  }

  // 创建解码器上下文
  codecCtx = avcodec_alloc_context3(codec);
  if (!codecCtx) {
    std::cout << "error: avcodec_alloc_context3() failed" << std::endl;
    goto fail;
  }

  // 使用视频流的编码参数设置解码器上下文
  ret = avcodec_parameters_to_context(
      codecCtx, formatCtx->streams[videoStream]->codecpar);
  if (ret < 0) {
    std::cout << "error: avcodec_parameters_to_context() failed" << std::endl;
    goto fail;
  }

  // 初始化解码器上下文
  ret = avcodec_open2(codecCtx, codec, nullptr);
  if (ret < 0) {
    std::cout << "error: avcodec_open2() failed" << std::endl;
    goto fail;
  }

  // 创建帧结构，RGB格式截图帧，用于保存视频帧转换后的结果
  frameRGB = av_frame_alloc();
  if (!frameRGB) {
    std::cout << "error: av_frame_alloc() failed" << std::endl;
    goto fail;
  }

  // 计算截图帧所需的内存大小
  int numBytes = av_image_get_buffer_size(AV_PIX_FMT_RGB24, codecCtx->width,
                                          codecCtx->height, 1);
  if (numBytes < 0) {
    std::cout << "error: av_image_get_buffer_size() failed" << std::endl;
    goto fail;
  }

  // 为截图帧分配内存
  buffer = (uint8_t *)av_malloc(numBytes);
  if (!buffer) {
    std::cout << "error: av_malloc() failed" << std::endl;
    goto fail;
  }

  // 将内存关联到截图帧frameRGB
  ret = av_image_fill_arrays(frameRGB->data, frameRGB->linesize, buffer,
                             AV_PIX_FMT_RGB24, codecCtx->width,
                             codecCtx->height, 1);
  if (ret < 0) {
    std::cout << "error: av_image_fill_arrays() failed" << std::endl;
    goto fail;
  }

  // 创建格式转换上下文
  // 用于执行视频帧到截图帧之间的像素格式转换，视频帧的像素格式一般为YUV
  SwsContext *swsCtx =
      sws_getContext(codecCtx->width, codecCtx->height, codecCtx->pix_fmt,
                     codecCtx->width, codecCtx->height, AV_PIX_FMT_RGB24,
                     SWS_BILINEAR, nullptr, nullptr, nullptr);

  // 创建帧结构，用于接收视频帧
  frame = av_frame_alloc();
  if (!frame) {
    std::cout << "error: av_frame_alloc() failed" << std::endl;
    goto fail;
  }

  // 循环读取 packet，媒体流被划分为多个 packet，将 packet 喂给解码器，
  // 然后从解码器获取解码后的帧，媒体流的一个 packet 包含一个或多个帧
  AVPacket packet;
  int frameFinished(0);
  int i(0);
  while (av_read_frame(formatCtx, &packet) >= 0) {
    // 判断是否是所需的媒体流
    if (packet.stream_index == videoStream) {
      // 将 packet 发送给解码器上下文
      ret = avcodec_send_packet(codecCtx, &packet);
      if (ret < 0) {
        std::cout << "error: avcodec_send_packet() failed" << std::endl;
        goto fail;
      }

      // 从解码器上下文中获取解码后的帧
      ret = avcodec_receive_frame(codecCtx, frame);
      if (ret == 0) {
        // 获取到一帧，执行像素格式转换
        sws_scale(swsCtx, frame->data, frame->linesize, 0, codecCtx->height,
                  frameRGB->data, frameRGB->linesize);

        // 将截图帧保存到文件，前5帧
        if (i < 5) {
          saveFrame(frameRGB, codecCtx->width, codecCtx->height, i);
          ++i;
        } else {
          break;
        }
      }
    }

    // 根据 av_read_frame() 函数的文档可知，函数成功时，
    // 需要调用 av_packet_unref() 释放内存
    av_packet_unref(&packet);
  }

fail:
  sws_freeContext(swsCtx);
  av_free(buffer);
  av_frame_free(&frameRGB);
  av_frame_free(&frame);
  avcodec_free_context(&codecCtx);
  avformat_close_input(&formatCtx);

  return 0;
}